// Copyright 1998-2016 Glenn McIntosh
// licensed under the GNU General Public Licence version 3
#pragma once

/** @file matrix.h
	Vector2D arithmetic and conversions.
	*/

// include files
#include <array>
#include "linear.h"

namespace math
{
//! Transpose a matrix.  @param a matrix to be transposed (can be rectangular) @return transposed matrix
template<size_t N, size_t M> Matrix<M, N> transpose(const Matrix<N, M> &a)
{
	int n{N}, m{M};
	Matrix<M, N> b;

	// for each row and column of a, transpose the elements
	for (int i = 0; i < n; ++i)
		for (int j = 0; j < m; ++j)
			b[j][i] = a[i][j];

	// return result
	return b;
}

//! Add two matrices. @param a first matrix @param b second matrix @return result matrix of a + b
template<size_t N> Array<N> operator+(Array<N> a, const Array<N> &b)
{
	int n{N};

	// for each column of a, add b
	for (int i = 0; i < n; ++i)
		a[i] += b[i];

	// return result
	return a;
}

//! Add two matrices. @param a first matrix @param b second matrix @return result matrix of a + b
template<size_t N, size_t M> Matrix<N, M> operator+(Matrix<N, M> a, const Matrix<N, M> &b)
{
	int n{N}, m{M};

	// for each row and column of a, add b
	for (int i = 0; i < n; ++i)
		for (int j = 0; j < m; ++j)
			a[i][j] += b[i][j];

	// return result
	return a;
}
//! Subtract two matrices. @param a first matrix @param b second matrix @return result matrix of a - b
template<size_t N> Array<N> operator-(Array<N> a, const Array<N> &b)
{
	int n{N};

	// for each column of a, subtract b
	for (int i = 0; i < n; ++i)
		a[i] -= b[i];

	// return result
	return a;
}
//! Subtract two matrices. @param a first matrix @param b second matrix @return result matrix of a - b
template<size_t N, size_t M> Matrix<N, M> operator-(Matrix<N, M> a, const Matrix<N, M> &b)
{
	int n{N}, m{M};

	// for each row and column of a, subtract b
	for (int i = 0; i < n; ++i)
		for (int j = 0; j < m; ++j)
			a[i][j] -= b[i][j];

	// return result
	return a;
}

//! Multiply two matrices. @param a first matrix @param b second matrix @return result matrix of a by b
template <size_t N, size_t M> Array<N> operator*(const Matrix<N, M> &a, const Array<M> &b)
{
	int n{N}, m{M};
	Array<N> r;

	// for each row of a, multiply out the elements
	for (int i = 0; i < n; ++i)
	{
		real s = 0.;
		for (int j = 0; j < m; ++j)
			s += a[i][j]*b[j];
		r[i] = s;
	}

	// return result
	return r;
}

//! Multiply two matrices. @param a first matrix @param b second matrix @return result matrix of a by b
template<size_t N, size_t M, size_t L> Matrix<N, L> operator*(const Matrix<N, M> &a, const Matrix<M, L> &b)
{
	int n{N}, m{M}, l{L};
	Matrix<N, L> r;

	// for each row of a and column of b, multiply out the elements
	for (int i = 0; i < n; ++i)
		for (int j = 0; j < l; ++j)
		{
			real s = 0.;
			for (int k = 0; k < m; ++k)
				s += a[i][k]*b[k][j];
			r[i][j] = s;
		}

	// return result
	return r;
}

//! Invert a matrix.  @param a matrix to be inverted @return inverted matrix
template<size_t N> Matrix<N, N> invert(Matrix<N, N> a)
{
	int n{N};
	real d;
	ArrayIndex<N> index;

	// decompose the matrix
	luDecompose(a, index, d);

	// create temporary
	Matrix<N, N> y{Array<N> {{}}};

	// solve for each column
	for (int j = 0; j < n; j++)
	{
		y[j][j] = 1.;
		luBackSubstitute(a, index, y[j]);
	}

	// transpose back into original matrix
	for (int j = 0; j < n; j++)
		for (int i = 0; i < n; i++)
			a[i][j] = y[j][i];

	// return result
	return a;
}

//! Invert a symmetric matrix.  
// @param a symmetric positive definite matrix to be inverted using Cholesky decomposition. @return inverted matrix
template<size_t N> Matrix<N, N> symmetricInvert(Matrix<N, N> a)
{
	int n{N};

	// convert source matrix into Cholesky factor
	choleskyDecompose(a);	

	// calculate inverse of Cholesky factor in upper triangle
	for (int i = 0; i < n; i++) 
		a[i][i] = 1 / a[i][i];
	for (int i = 0; i < n; i++) 
		for (int j = i+1; j < n; j++) 
		{
			a[j][i] *= a[i][i];
			for (int k = i+1; k < j; k++) 
				a[j][i] += a[j][k] * a[k][i];
			a[j][i] *= -a[j][j];
		}
	
	// calculate (Cholesky factor inverted)' * (Cholesky factor inverted)
	for (int i = 0; i < n; i++) 
	{
		a[i][i] *= a[i][i];
		for (int k = i+1; k < n; k++) 
			a[i][i] += a[k][i] * a[k][i];
		for (int j = i+1; j < n; j++) 
		{
			a[j][i] *= a[j][j];
			for (int k = j+1; k < n; k++) 
				a[j][i] += a[k][j] * a[k][i];
			a[i][j] = a[j][i];
		}
	}
	return a;
}
}
